import requests
import json

'''
@input
{
    "data": [{"title":title,"content":content,"dataid":dataid},
             {"title":title,"content":content,"dataid":dataid}
             ...      
             ]
}

'''

class TextToVector():
    def __init__(self, host='10.39.14.233', port=3451):
        self.__host = host
        self.__port = port
        self.headers = {'Content-Type': 'application/json',"connection":"close"}

    def get_vec(self, text, text_other='', dataid='_'):
        assert isinstance(text, str)

        data = json.dumps({"data":[{"title":text,
                                    "content":text_other if text_other else text,
                                    "dataid":dataid}]})
        try:
            response = requests.post(url='http://{}:{}/bert_dssm/'.format(self.__host, self.__port), 
                            headers=self.headers, data=data)
            dssm_vectors = json.loads(response.json()['ret'])
            return dssm_vectors[dataid]
        except Exception as e:
            print(e)
            return []
    
    def __call__(self, text):
        return self.get_vec(text)


if __name__ == "__main__":
    text = '《少年的你》代表香港角逐奥斯卡最佳国际影片'
    vec = TextToVector()(text)
    print(vec)
